# https://blog.csdn.net/bananapai/article/details/145736300
from torchvision.datasets import MNIST

# 下载完毕之后我们可以手动设置download为False
# ds=MNIST(root='./data',train=True,download=True)
ds = MNIST(root='./data', train=True, download=False)

# print(len(ds))

import torch
from torch import nn


class AlexNet(nn.Module):
    def __init__(self):
        super().__init__()
        # super(AlexNet,self).__init__()

        self.features = nn.Sequential(
            # 96个卷积核，卷积核为11，stride为11，padding为2,output为 [-1, 96, 55, 55],The number of parameters=1*96*11*11+96=11,712
            nn.Conv2d(1, 96, 11, 4, 2),
            # nn.Conv2d(1, 96, 11, 4, 3),
            nn.ReLU(),
            # output=[-1, 96, 27, 27]
            nn.MaxPool2d(3, stride=2),

            # 256个卷积核，卷积核为5，stride=1,padding=2,output=[-1, 256, 27, 27] ,The number of parameters=96*256*5*5+256=614,656
            nn.Conv2d(96, 256, 5, padding=2),
            nn.ReLU(),
            # output=[-1, 256, 13, 13]
            nn.MaxPool2d(3, 2),

            # 384个卷积核，卷积核为3，stride=1,padding=1,no maxpool,output=[-1, 384, 13, 13],The number of parameters=256*384*3*3+384=885,120
            nn.Conv2d(256, 384, 3, padding=1),
            nn.ReLU(),

            # 384 convolution kernels,the size of convolution kernels is 3.stride=1,padding=1,no maxpool,
            # output=[-1, 384, 13, 13], The number of parameters=384*384*3*3+384=1,327,488
            nn.Conv2d(384, 384, 3, padding=1),
            nn.ReLU(),

            # 256 convolution kernels,the size of covolution kernels is 3,stride=1,padding=1,output=[-1, 256, 13, 13]
            # The number of parameters=384*256*3*3+256=884,992
            nn.Conv2d(384, 256, 3, padding=1),
            nn.ReLU(),
            # output=[-1, 256, 6, 6]
            nn.MaxPool2d(3, 2),

            # the probability=0.5
            nn.Dropout(p=0.5),
            # output=[-1, 9216]
            nn.Flatten(start_dim=1)
        )

        self.classifier = nn.Sequential(
            # input neurons=9216,output neurons=4096,output= [-1, 4096],The number of parameters=9216*4096+4096=37,752,832
            nn.Linear(9216, 4096),
            nn.ReLU(),
            # input neurons and output neurons are both 4096,output= [-1, 4096],The number of parameters=4096*4096+4096=16,781,312
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(p=0.3),
            # input neurons=4096,output neurons=1000,output=[-1, 10] ,The number of parameters=4096*10+10=40970
            nn.Linear(4096, 10),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


# import torch
## for test
model = AlexNet()
from torchsummary import summary

summary(model, (1, 224, 224))
for p in model.parameters():
    print(p.shape)

# x = torch.rand(10, 1, 224, 224)
# y = model(x)
# print(y.shape)
print("alexnet2mnist_model", 444)